"""
任务管理API路由
"""
import logging
from typing import List, Optional
from datetime import datetime, timedelta
from fastapi import APIRouter, Depends, HTTPException, Query, Path
from sqlalchemy.orm import Session
from sqlalchemy import desc
from pydantic import BaseModel, Field

from modules.models.base import get_db
from modules.models.stock import CollectionTask
from modules.tasks.celery_app import app
from modules.api.auth import get_current_active_user, User

# 配置日志
logger = logging.getLogger("api.tasks")

# 创建路由
router = APIRouter()


# 模型定义
class TaskBase(BaseModel):
    """任务基础信息模型"""
    task_id: Optional[str] = None
    task_type: str
    security_code: Optional[str] = None
    status: str
    create_time: datetime
    update_time: datetime
    
    class Config:
        orm_mode = True


class TaskDetail(TaskBase):
    """任务详细信息模型"""
    start_date: Optional[datetime] = None
    end_date: Optional[datetime] = None
    retry_count: int
    error_message: Optional[str] = None
    result_count: int
    
    class Config:
        orm_mode = True


class TaskResponse(BaseModel):
    """任务操作响应模型"""
    success: bool
    message: str
    task_id: Optional[str] = None


# API路由
@router.get("/list", response_model=List[TaskBase], summary="获取任务列表")
async def get_task_list(
    task_type: Optional[str] = Query(None, description="任务类型"),
    status: Optional[str] = Query(None, description="任务状态"),
    security_code: Optional[str] = Query(None, description="证券代码"),
    limit: int = Query(50, description="返回结果限制数量", ge=1, le=1000),
    offset: int = Query(0, description="结果偏移量", ge=0),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_active_user)
):
    """
    获取任务列表
    
    - **task_type**: 任务类型，可选
    - **status**: 任务状态，可选
    - **security_code**: 证券代码，可选
    - **limit**: 返回结果数量限制，默认50，最大1000
    - **offset**: 结果偏移量，默认0
    """
    try:
        query = db.query(CollectionTask)
        
        if task_type:
            query = query.filter(CollectionTask.task_type == task_type)
        
        if status:
            query = query.filter(CollectionTask.status == status)
        
        if security_code:
            query = query.filter(CollectionTask.security_code == security_code)
        
        total = query.count()
        tasks = query.order_by(desc(CollectionTask.create_time)).offset(offset).limit(limit).all()
        
        return tasks
    
    except Exception as e:
        logger.error(f"获取任务列表失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"获取任务列表失败: {str(e)}")


@router.get("/{task_id}", response_model=TaskDetail, summary="获取任务详情")
async def get_task_detail(
    task_id: str = Path(..., description="任务ID"),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_active_user)
):
    """
    获取任务详情
    
    - **task_id**: 任务ID
    """
    try:
        # 查询数据库中的任务记录
        task = db.query(CollectionTask).filter(CollectionTask.task_id == task_id).first()
        
        if not task:
            # 如果数据库中没有记录，尝试从Celery查询任务状态
            celery_task = app.AsyncResult(task_id)
            if celery_task.state == 'PENDING':
                raise HTTPException(status_code=404, detail=f"任务 {task_id} 不存在")
            
            # 创建一个临时任务对象
            temp_task = CollectionTask(
                task_id=task_id,
                task_type="unknown",
                status=celery_task.state.lower(),
                create_time=datetime.now(),
                update_time=datetime.now(),
                retry_count=0,
                result_count=0
            )
            
            if celery_task.failed():
                temp_task.error_message = str(celery_task.result)
            
            return temp_task
        
        return task
    
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取任务 {task_id} 详情失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"获取任务详情失败: {str(e)}")


@router.post("/{task_id}/revoke", response_model=TaskResponse, summary="取消任务")
async def revoke_task(
    task_id: str = Path(..., description="任务ID"),
    terminate: bool = Query(False, description="是否终止任务"),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_active_user)
):
    """
    取消任务
    
    - **task_id**: 任务ID
    - **terminate**: 是否终止任务，默认为False（等待任务结束），True表示立即终止
    """
    try:
        # 检查任务是否存在
        task = db.query(CollectionTask).filter(CollectionTask.task_id == task_id).first()
        
        # 取消Celery任务
        app.control.revoke(task_id, terminate=terminate)
        
        # 如果任务存在于数据库中，更新其状态
        if task:
            task.status = 'revoked'
            task.update_time = datetime.now()
            db.commit()
        
        return {
            "success": True,
            "message": f"任务 {task_id} 已取消" + (" 并立即终止" if terminate else ""),
            "task_id": task_id
        }
    
    except Exception as e:
        logger.error(f"取消任务 {task_id} 失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"取消任务失败: {str(e)}")


@router.post("/{task_id}/retry", response_model=TaskResponse, summary="重试任务")
async def retry_task(
    task_id: str = Path(..., description="任务ID"),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_active_user)
):
    """
    重试已失败的任务
    
    - **task_id**: 任务ID
    """
    try:
        # 检查任务是否存在
        task = db.query(CollectionTask).filter(CollectionTask.task_id == task_id).first()
        
        if not task:
            raise HTTPException(status_code=404, detail=f"任务 {task_id} 不存在")
        
        if task.status != 'failed':
            raise HTTPException(status_code=400, detail=f"只能重试失败的任务，当前状态: {task.status}")
        
        # 创建新的任务
        from modules.tasks.stock_tasks import collect_stock_data
        
        # 根据任务类型决定调用哪个任务函数
        if task.task_type == 'stock_daily':
            start_date = task.start_date.strftime('%Y%m%d') if task.start_date else None
            end_date = task.end_date.strftime('%Y%m%d') if task.end_date else None
            
            new_task = collect_stock_data.delay(task.security_code, start_date, end_date)
            
            # 更新原任务状态
            task.status = 'retried'
            task.update_time = datetime.now()
            db.commit()
            
            return {
                "success": True,
                "message": f"任务 {task_id} 已重新提交",
                "task_id": new_task.id
            }
        else:
            raise HTTPException(status_code=400, detail=f"不支持重试的任务类型: {task.task_type}")
    
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"重试任务 {task_id} 失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"重试任务失败: {str(e)}")


@router.post("/batch_collect", response_model=TaskResponse, summary="批量采集数据")
async def batch_collect_data(
    task_type: str = Query(..., description="任务类型: stock_daily/index_daily/etf_daily"),
    days: int = Query(30, description="采集天数", ge=1, le=365),
    security_codes: Optional[List[str]] = Query(None, description="证券代码列表"),
    current_user: User = Depends(get_current_active_user)
):
    """
    批量采集数据
    
    - **task_type**: 任务类型
    - **days**: 采集天数，默认30天，最大365天
    - **security_codes**: 证券代码列表，可选，不提供则采集全部
    """
    try:
        # 计算日期范围
        end_date = datetime.now().strftime('%Y%m%d')
        start_date = (datetime.now() - timedelta(days=days)).strftime('%Y%m%d')
        
        # 根据任务类型决定调用哪个模块
        if task_type == 'stock_daily':
            from modules.tasks.stock_tasks import collect_stock_data, schedule_daily_stock_update
            
            if security_codes:
                # 创建任务组
                from celery import group
                
                tasks = []
                for code in security_codes:
                    tasks.append(collect_stock_data.s(code, start_date, end_date))
                
                result = group(tasks).apply_async()
                
                return {
                    "success": True,
                    "message": f"已提交 {len(security_codes)} 只股票的数据采集任务",
                    "task_id": str(result.id)
                }
            else:
                # 采集全部股票
                result = schedule_daily_stock_update.delay()
                
                return {
                    "success": True,
                    "message": "已提交全部股票的数据采集任务",
                    "task_id": result.id
                }
        else:
            raise HTTPException(status_code=400, detail=f"不支持的任务类型: {task_type}")
    
    except Exception as e:
        logger.error(f"批量提交任务失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"批量提交任务失败: {str(e)}")


@router.get("/statistics", summary="获取任务统计信息")
async def get_task_statistics(
    days: int = Query(7, description="统计天数", ge=1, le=30),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_active_user)
):
    """
    获取任务统计信息
    
    - **days**: 统计天数，默认7天，最大30天
    """
    try:
        # 计算起始日期
        start_date = datetime.now() - timedelta(days=days)
        
        # 获取任务类型统计
        task_types = db.query(
            CollectionTask.task_type,
            func.count(CollectionTask.id).label('count')
        ).filter(
            CollectionTask.create_time >= start_date
        ).group_by(
            CollectionTask.task_type
        ).all()
        
        # 获取任务状态统计
        task_status = db.query(
            CollectionTask.status,
            func.count(CollectionTask.id).label('count')
        ).filter(
            CollectionTask.create_time >= start_date
        ).group_by(
            CollectionTask.status
        ).all()
        
        # 获取每日任务数量
        daily_tasks = db.query(
            func.date(CollectionTask.create_time).label('date'),
            func.count(CollectionTask.id).label('count')
        ).filter(
            CollectionTask.create_time >= start_date
        ).group_by(
            func.date(CollectionTask.create_time)
        ).all()
        
        # 构建返回结果
        result = {
            "task_types": {t.task_type: t.count for t in task_types},
            "task_status": {s.status: s.count for s in task_status},
            "daily_tasks": {str(d.date): d.count for d in daily_tasks},
            "total_tasks": sum(t.count for t in task_types)
        }
        
        return result
    
    except Exception as e:
        logger.error(f"获取任务统计信息失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"获取任务统计信息失败: {str(e)}")